#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
频段功率计算工具

用于从时域信号计算各频段的功率值，支持正确的信号处理流程：
时域信号 -> FFT -> 频段功率计算
"""

import numpy as np
from scipy import signal
from typing import Dict, List, Optional, Tuple
import logging

logger = logging.getLogger(__name__)


class FrequencyBandCalculator:
    """
    频段功率计算器
    
    实现正确的频段功率计算流程：
    1. 接收时域信号数据
    2. 应用窗函数减少频谱泄漏
    3. 进行FFT变换
    4. 计算各频段的功率谱密度
    5. 返回标准化的频段功率值
    """
    
    def __init__(self, sampling_rate: float = 512.0, window_size: int = 256):
        """
        初始化频段功率计算器
        
        Args:
            sampling_rate: 采样率 (Hz)
            window_size: FFT窗口大小，影响频率分辨率
        """
        self.sampling_rate = sampling_rate
        self.window_size = window_size
        
        # 定义8频域脑电频段（支持更精细分析）
        self.frequency_bands_8 = {
            'delta': (0.5, 4.0),        # Delta波段
            'theta': (4.0, 8.0),        # Theta波段
            'low_alpha': (8.0, 10.0),   # Low Alpha波段
            'high_alpha': (10.0, 13.0), # High Alpha波段
            'low_beta': (13.0, 20.0),   # Low Beta波段
            'high_beta': (20.0, 30.0),  # High Beta波段
            'low_gamma': (30.0, 50.0),  # Low Gamma波段
            'mid_gamma': (50.0, 100.0)  # Mid Gamma波段
        }
        
        # 保留5频域定义（向后兼容）
        self.frequency_bands_5 = {
            'delta': (0.5, 4.0),    # Delta波段
            'theta': (4.0, 8.0),    # Theta波段
            'alpha': (8.0, 13.0),   # Alpha波段
            'beta': (13.0, 30.0),   # Beta波段
            'gamma': (30.0, 100.0)  # Gamma波段
        }
        
        # 默认使用8频域
        self.frequency_bands = self.frequency_bands_8
        
        # 计算频率分辨率
        self.freq_resolution = sampling_rate / window_size
        
        logger.info(f"频段功率计算器初始化完成:")
        logger.info(f"  - 采样率: {sampling_rate}Hz")
        logger.info(f"  - 窗口大小: {window_size}")
        logger.info(f"  - 频率分辨率: {self.freq_resolution:.2f}Hz")
    
    def calculate_band_powers_from_signal(self, time_signal: np.ndarray) -> Dict[str, int]:
        """
        从时域信号计算各频段功率
        
        Args:
            time_signal: 时域信号数组
            
        Returns:
            Dict[str, int]: 各频段功率值字典
        """
        try:
            # 确保信号长度足够
            if len(time_signal) < self.window_size:
                # 如果信号太短，进行零填充
                padded_signal = np.zeros(self.window_size)
                padded_signal[:len(time_signal)] = time_signal
                time_signal = padded_signal
                logger.debug(f"信号长度不足，已零填充至{self.window_size}")
            
            # 取最后window_size个样本进行分析
            if len(time_signal) > self.window_size:
                time_signal = time_signal[-self.window_size:]
            
            # 去除直流分量
            time_signal = time_signal - np.mean(time_signal)
            
            # 计算功率谱密度（移除双重窗函数问题）
            # signal.periodogram内部已经应用汉宁窗，无需手动再次应用
            freqs, psd = signal.periodogram(
                time_signal, 
                fs=self.sampling_rate, 
                window='hann',
                scaling='density'
            )
            
            # 计算各频段功率
            band_powers = {}
            for band_name, (low_freq, high_freq) in self.frequency_bands.items():
                power = self._calculate_band_power(freqs, psd, low_freq, high_freq)
                # 转换为整数并确保非负
                band_powers[band_name] = max(0, int(power))
            
            logger.debug(f"频段功率计算完成: {band_powers}")
            return band_powers
            
        except Exception as e:
            logger.error(f"频段功率计算失败: {e}")
            # 返回默认8频域值
            return {
                'delta': 0,
                'theta': 0,
                'low_alpha': 0,
                'high_alpha': 0,
                'low_beta': 0,
                'high_beta': 0,
                'low_gamma': 0,
                'mid_gamma': 0
            }
    
    def calculate_band_powers_from_buffer(self, signal_buffer: List[float]) -> Dict[str, int]:
        """
        从信号缓冲区计算各频段功率
        
        Args:
            signal_buffer: 信号缓冲区列表
            
        Returns:
            Dict[str, int]: 各频段功率值字典
        """
        if not signal_buffer:
            return {
                'delta': 0,
                'theta': 0,
                'low_alpha': 0,
                'high_alpha': 0,
                'low_beta': 0,
                'high_beta': 0,
                'low_gamma': 0,
                'mid_gamma': 0
            }
        
        # 转换为numpy数组
        time_signal = np.array(signal_buffer, dtype=float)
        return self.calculate_band_powers_from_signal(time_signal)
    
    def _calculate_band_power(self, freqs: np.ndarray, psd: np.ndarray, 
                             low_freq: float, high_freq: float) -> float:
        """
        计算指定频段的功率
        
        Args:
            freqs: 频率数组
            psd: 功率谱密度数组
            low_freq: 频段下限
            high_freq: 频段上限
            
        Returns:
            float: 频段功率值
        """
        # 找到频段范围内的索引
        band_mask = (freqs >= low_freq) & (freqs <= high_freq)
        
        if not np.any(band_mask):
            logger.warning(f"频段 {low_freq}-{high_freq}Hz 没有对应的频率点")
            return 0.0
        
        # 计算频段内的总功率（积分）
        band_freqs = freqs[band_mask]
        band_psd = psd[band_mask]
        
        # 使用梯形积分计算功率
        power = np.trapz(band_psd, band_freqs)
        
        # 归一化到合理范围（类似NeuroSky的输出范围）
        # 优化：调整归一化系数到合理范围，避免数值异常放大
        # 根据实际测试结果，使用10000000作为归一化系数以匹配原始数据范围
        normalized_power = power * 10000000  # 优化后的归一化系数
        
        return max(0.0, normalized_power)
    
    def get_frequency_resolution(self) -> float:
        """
        获取当前的频率分辨率
        
        Returns:
            float: 频率分辨率 (Hz)
        """
        return self.freq_resolution
    
    def get_supported_bands(self) -> Dict[str, Tuple[float, float]]:
        """
        获取支持的频段定义
        
        Returns:
            Dict[str, Tuple[float, float]]: 频段名称到频率范围的映射
        """
        return self.frequency_bands.copy()
    
    def calculate_band_powers_8bands(self, time_signal: np.ndarray) -> Dict[str, int]:
        """
        计算8频域功率（明确指定8频域模式）
        
        Args:
            time_signal: 时域信号数组
            
        Returns:
            Dict[str, int]: 8频域功率值字典
        """
        # 临时切换到8频域模式
        original_bands = self.frequency_bands
        self.frequency_bands = self.frequency_bands_8
        
        try:
            result = self.calculate_band_powers_from_signal(time_signal)
            return result
        finally:
            # 恢复原始频段设置
            self.frequency_bands = original_bands
    
    def calculate_band_powers_5bands(self, time_signal: np.ndarray) -> Dict[str, int]:
        """
        计算5频域功率（向后兼容模式）
        
        Args:
            time_signal: 时域信号数组
            
        Returns:
            Dict[str, int]: 5频域功率值字典
        """
        # 临时切换到5频域模式
        original_bands = self.frequency_bands
        self.frequency_bands = self.frequency_bands_5
        
        try:
            result = self.calculate_band_powers_from_signal(time_signal)
            return result
        finally:
            # 恢复原始频段设置
            self.frequency_bands = original_bands
    
    def update_band_definition(self, band_name: str, low_freq: float, high_freq: float):
        """
        更新频段定义
        
        Args:
            band_name: 频段名称
            low_freq: 频段下限
            high_freq: 频段上限
        """
        if low_freq >= high_freq:
            raise ValueError(f"频段下限({low_freq})必须小于上限({high_freq})")
        
        if high_freq > self.sampling_rate / 2:
            raise ValueError(f"频段上限({high_freq})不能超过奈奎斯特频率({self.sampling_rate/2})")
        
        self.frequency_bands[band_name] = (low_freq, high_freq)
        logger.info(f"频段定义已更新: {band_name} = {low_freq}-{high_freq}Hz")


# 便捷函数
def calculate_band_powers(time_signal: np.ndarray, sampling_rate: float = 512.0) -> Dict[str, int]:
    """
    便捷函数：从时域信号计算频段功率
    
    Args:
        time_signal: 时域信号
        sampling_rate: 采样率
        
    Returns:
        Dict[str, int]: 频段功率字典
    """
    calculator = FrequencyBandCalculator(sampling_rate=sampling_rate)
    return calculator.calculate_band_powers_from_signal(time_signal)


def create_band_calculator(sampling_rate: float = 512.0, window_size: int = 256) -> FrequencyBandCalculator:
    """
    便捷函数：创建频段功率计算器
    
    Args:
        sampling_rate: 采样率
        window_size: FFT窗口大小
        
    Returns:
        FrequencyBandCalculator: 频段功率计算器实例
    """
    return FrequencyBandCalculator(sampling_rate=sampling_rate, window_size=window_size)